import gymnasium
from gymnasium import spaces
from odps import ODPS
import time
from sklearn.preprocessing import LabelEncoder
import numpy as np
import copy
import pandas as pd
from sklearn.cluster import KMeans
from tensorboardX import SummaryWriter
import torch
import torch.nn as nn
import torch.nn.functional as F

bidincome_total = 0
transaction = 0
theoretical_total = 0
episode_reward_total = 0

class DictStateEmbedding(nn.Module):
    def __init__(self, categorical_feature_dimensions, embedding_dims, continuous_features_dim, output_size):
        """
        :param categorical_feature_dimensions: 字典，键为类别特征的名称，值为每个类别特征的可能最大值加1
        :param embedding_dims: 字典，键为类别特征的名称，值为对应的嵌入维度
        :param continuous_features_dim: 连续特征的数量
        :param output_size: 网络输出层的大小
        """
        super(DictStateEmbedding, self).__init__()
        self.embeddings = nn.ModuleDict({
            key: nn.Embedding(num_embeddings=value, embedding_dim=embedding_dims[key])
            for key, value in categorical_feature_dimensions.items()
        })
        # self.fc_continuous = nn.Linear(continuous_features_dim, 64)  # 处理连续特征的全连接层
        total_embedding_dim = sum(embedding_dims.values()) + continuous_features_dim  # 所有嵌入维度之和 + 连续特征处理后的维度

        # 定义后续的全连接层
        self.fc1 = nn.Linear(total_embedding_dim, 128)
        self.fc2 = nn.Linear(128, output_size)

    def forward(self, x):
        """
        :param x: 字典类型的输入，包含类别特征和连续特征
        """
        embeddings = []
        for key, module in self.embeddings.items():
            embeddings.append(module(x[key]))
        embeddings = torch.cat(embeddings, dim=-1)

        x = F.relu(self.fc1(embeddings))
        x = self.fc2(x)
        return x


        

EMB_MAX = 1
output_size = 8  # 输出向量的维度
model = DictStateEmbedding(categorical_feature_dimensions, embedding_dims, continuous_features_dim, output_size)



class MyEnv(gymnasium.Env):
    
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 4}

    def __init__(self):
        self.state = {}
        self.episode_length = np.array([data_num], dtype='int')
        self.expected_income = np.array([data_num * 300], dtype='float32')
        self.action_space = spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32)
        self.data_counter = 0
        self.episode_reward = 0
        self.n_episode = 0
        self.bidincome = 0
        self.transaction = 0
        self.df = df_all
        self.theoretical = 0
        self.bidincome_total = 0
        self.theoretical_total = 0
        self.episode_reward_total = 0 
        self.rationality = 0 
        self.all_bidincome_totel = 0
        self.all_transaction = 0
        self.all_theoretical_total = 0
        self.all_rationality = 0   
        self.observation_space = spaces.Box(np.array([0]+ [-np.inf]*9), np.array([np.inf]*10), dtype=np.float32)

    def reset(self, seed=None, options=None):
        # print("_____________reset________________")
        self.data_counter = 0
        self.bidincome = 0
        self.theoretical = 0
        if self.n_episode != 0:
            writer.add_scalar('episode_reward', self.episode_reward, self.n_episode)
        self.n_episode += 1
        self.episode_reward = 0
        embedded_flowvec_list, _ = self.df

        info = {}
        self.state["remaining_rounds"] = copy.deepcopy(self.episode_length)
        self.state["remaining_income"] = copy.deepcopy(self.expected_income)
        self.state['flow_vec'] = model(embedded_flowvec_list[self.data_counter]).detach().numpy()
        state_vec = []
        for value in self.state.values():
            for _value in value:
                state_vec.append(_value)
        state_vec[0] = float(state_vec[0] / 1024) * EMB_MAX
        state_vec[1] = float(state_vec[1] / 307200) * EMB_MAX
        obs = np.array(state_vec, dtype='float32')
        
        return obs, info
    
    def step(self, action):
        embedded_flowvec_list, bidprice_list = self.df
        bid_prices = np.array(bidprice_list[self.data_counter], dtype='float32')
        
        if self.data_counter < self.episode_length - 1:
            self.data_counter += 1
        else:
            self.data_counter = 0
        
        action = action * 37.5 + 262.5         
       
        
        self.state['remaining_rounds'] -= 1

        self.state['flow_vec'] = model(embedded_flowvec_list[self.data_counter]).detach().numpy()
        


        if action[0] <= bid_prices:
            self.transaction +=1
            self.all_transaction += 1
            reward = action[0]
            bidincome = float(action[0])
            self.state['remaining_income'] -= float(action[0])
        else:
            reward = -float(action[0] - bid_prices) 
            bidincome = 0
     
        if abs(action[0] - bid_prices) <= 15:
            self.rationality += 1  
            self.all_rationality += 1

        self.theoretical += bid_prices
        self.episode_reward += reward
        self.bidincome += bidincome
        self.bidincome_total += bidincome
        self.episode_reward_total += reward
        self.theoretical_total += bid_prices


        self.all_bidincome_totel += bidincome
        self.all_theoretical_total += bid_prices
        

        info = {}
        if self.state['remaining_rounds'] == 0:
            done = True
        else:
            done = False      
        # print(self.state, reward, done, info)
        
        state_vec = []
        for value in self.state.values():
            for _value in value:
                state_vec.append(_value)
        
        state_vec[0] = float(state_vec[0] / 1024) * EMB_MAX
        state_vec[1] = float(state_vec[1] / 307200) * EMB_MAX
        obs = np.array(state_vec, dtype='float32')

        return obs, reward, done, False, info
    def print(self):
        data_epside_num = 200
        if (self.data_counter +1) % data_epside_num == 0 and self.data_counter != 0:
        # if self.data_counter  == (data_num - 1):
            if self.transaction == 0:
                self.transaction = 1
            if self.all_transaction == 0:
                self.all_transaction = 1
            average_income = self.bidincome_total / data_epside_num
            eva_reward = self.episode_reward_total / data_epside_num
            average_transactionprice = self.bidincome_total / self.transaction
            difference = (self.theoretical_total - self.bidincome_total)/data_epside_num
            theoretical_price = self.theoretical_total/data_epside_num
            rationality_total = self.rationality/data_epside_num
            turnover_rate = self.transaction / data_epside_num
            print("当前数据量：", self.data_counter + 1, end = " ")
            # print("平均奖励: ", '%.4f'% eva_reward, end = " ")
            print("平均出价: ", '%.2f'% theoretical_price, end = " ")
            print("平均收益: ", '%.2f'% average_income, end = " ")
            print("差值: ", '%.2f'% difference, end = " ")
            print("平均成交价: ", '%.2f'% average_transactionprice, end = " ")
            print("成交率: ", '%.2f'% turnover_rate, end = " ")
            print("合理度："'%.2f'% rationality_total)
            self.bidincome_total = 0
            self.transaction = 0
            self.theoretical_total = 0
            self. episode_reward_total = 0
            self.rationality = 0
            return theoretical_price, turnover_rate, rationality_total
        if self.data_counter == (data_num - 1):
            a = self.all_theoretical_total/data_num
            b = self.all_bidincome_totel/data_num
            c = (self.all_theoretical_total - self.all_bidincome_totel)/data_num
            d = self.all_bidincome_totel/self.all_transaction
            e = self.all_transaction/data_num
            f = self.all_rationality/data_num
            print("总数据量：", data_num, end=" ")
            print("总平均出价：", '%.2f'% a, end=" ")
            print("总平均收益：", '%.2f'% b, end=" ")
            print("总差值: ", '%.2f'% c, end = " ")
            print("总平均成交价: ", '%.2f'% d, end = " ")
            print("总成交率: ", '%.2f'% e, end = " ")
            print("总合理度："'%.2f'% f)
            return  None
        else:
            return  None
    def close(self):
        pass
    
    def render(self):
        pass

    def seed(self, seed=None):
        pass
